#include <iostream>
#include <vector>
#include <string>
#include <unistd.h>
#include <cstdlib>
#include <wait.h>
#include "task.hpp"
using namespace std;

enum
{
    USAGE_ERROR = 1,
    ARG_ERROR,
    PIPE_ERROR
};

class Channel
{
public:
    Channel(int wfd, int subprocess_id, string cname)
        : _wfd(wfd), _subprocess_id(subprocess_id), _cname(cname)
    {
    }
    int Wfd()
    {
        return _wfd;
    }
    int id()
    {
        return _subprocess_id;
    }
    string &name()
    {
        return _cname;
    }
    void debug()
    {
        cout << "wfd:" << _wfd << ", sub process id:" << _subprocess_id << ", name:" << _cname << endl;
    }
    void Close() { close(_wfd); } //关闭管道写端，确保子进程在读取到 EOF 后退出任务监听循环

private:
    int _wfd;
    int _subprocess_id;
    string _cname;
};
class ProcessPool
{
public:
    ProcessPool(int subprocess_num)
        : _sub_process_num(subprocess_num)
    {
    }
    ~ProcessPool()
    {}
    void Debug()
    {
        for (auto &e : _channels)
            e.debug();
    }
    int CreateProcess(work_t work) // 回调函数
    {
        vector<int> fds;
        _channels.reserve(_sub_process_num);
        for (int i = 0; i < _sub_process_num; i++)
        {
            // sleep(1);
            int pipefd[2]{0};
            int n = pipe(pipefd);
            if (n < 0)
            {
                return PIPE_ERROR;
            }
            fds.push_back(pipefd[1]);

            // 有一个bug,除了0号管道，所有的管道的父进程和子进程不是一一对应的 在KillAll函数中复现这个bug
            pid_t id = fork();
            if (id == 0)
            {
                // child - r
                for (auto wfd : fds)
                    close(wfd);
                dup2(pipefd[0], 0); // 重定向标准输入
                work(pipefd[0]);
                exit(0);
            }
            // father - w
            close(pipefd[0]);
            string cname = "Channel-" + to_string(i);
            _channels.emplace_back(pipefd[1], id, cname);
        }
        return 0;
    }
    int NextChannel()
    {
        static int cnt = 0;
        cnt %= _sub_process_num;
        int ret = cnt;
        cnt++;
        return ret;
    }
    void SendTask(int index, uint32_t command_code)
    {
        // 向指定管道写入任务码
        cout << "Send task:" << command_code << " to " << _channels[index].name() << ", sub process id is " << _channels[index].id() << endl;
        write(_channels[index].Wfd(), &command_code, sizeof(command_code));
    }
    void ProcessControl(int cnt)
    {
        while (cnt--)
        {
            // a.选择一个管道
            int channel = NextChannel();
            // b.选择一个任务
            uint32_t command_code = NextTask();
            // a.发送任务
            SendTask(channel, command_code);
            sleep(1);
        }
        cout << "task run done \n";
    }
    // 让子进程退出，只需要关闭管道的wfd即可
    // 解决bug办法:
    /*1.从最后一个创建的子进程开始等待
      2.所有进程关闭后在等待
      3.使用vector记录父进程前面所有的wfd，创建新的子进程时close（vector的wfd）
    */
    void KillAll()
    {
        // 退出一个子进程，回收一个资源
        for (auto e : _channels)
        {
            e.Close();  // 关闭管道的写端，读端进程(子进程)自动退出
            pid_t pid = e.id();
            pid_t wid = waitpid(pid, nullptr, 0); // 可能会出现bug，父进程阻塞等待子进程退出，但是子进程一直在read阻塞
            if (pid == wid)
            {
                cout << "wait success! sub pid is:" << pid << endl;
            }
            cout  << e.name()  << " close" << ", sub_process id: " << e.id() << "\n" << endl;
            // sleep(1);
        }
    }
    // 回收子进程资源
    void Wait()
    {
        for (auto &e : _channels)
        {
            pid_t pid = e.id();
            pid_t wid = waitpid(pid, nullptr, 0);
            if (pid == wid)
            {
                cout << "wait success! id is:" << pid << endl;
            }
        }
    }

private:
    int _sub_process_num;
    std::vector<Channel> _channels;
};
int main(int argc, char *argv[])
{
    // 0.获取命令行参数
    srand((uint64_t)time(nullptr));
    int subprocess_num = 0;
    if (argc != 2)
    {
        std::cout << "Usage: ./processpool sub_process_num" << std::endl;
        return USAGE_ERROR;
    }
    subprocess_num = std::stoi(argv[1]);
    if (subprocess_num < 0)
    {
        cout << "sub process should >= 0";
        return ARG_ERROR;
    }

    // 1.创建通信信道和子进程
    ProcessPool *processpool_ptr = new ProcessPool(subprocess_num);
    processpool_ptr->CreateProcess(Worker);

    //  2.控制子进程
    processpool_ptr->ProcessControl(5);

    //  3.回收子进程
    // a.让所有的子进程退出
    processpool_ptr->KillAll();
    // b.子进程资源回收 wait/waitpid
    // processpool_ptr->Wait();
    delete processpool_ptr;
    return 0;
}